import os

import cv2
import numpy as np
import torch
import torchvision
from ultralytics import YOLO

if __name__ == '__main__':
    # 加载训练好的模型，用YOLO加载
    model = YOLO(r'./runs/pose/train7/weights/best.pt')

    # 将cv2读取的图片转为张量(归一化到[0,1])
    transformer = torchvision.transforms.ToTensor()
    img = cv2.imread(r'./imgs/plate2.jpg')
    H,W = img.shape[:2]
    imgTensor = transformer(img)
    # 注意模型传入需要[Batch, C,H,W]形状的张量，所以reshape一下
    imgTensor = imgTensor.reshape(-1,3,H, W)

    # 前向传播
    model.eval()  # 注意YOLO的train似乎被重写了？不需要torch自建模型那样喂数据训练
    # result是一个列表，每一个元素对应一张图片的模型输出结果，结果解析打断点看或者见main.py
    result = model(imgTensor)
    print(result)
    result[0].show()